#include "AnaDef.h"
#include "XGLUtils.h"
#include "dEdxGausN.h"
#include "style.h"

Float_t ftpcsig = -999;
Float_t fpt = -999;
Float_t fmom = -999;

Float_t feta = -999;  
Float_t fv0cen = -999;
Int_t fcharge = -999;

//Float_t ftofmass = -999;
Float_t ftofpion = -999;
Float_t ftofkaon = -999;
Float_t ftofproton = -999;
Float_t ftofelectron = -999;

Int_t fifilled = -999;

Int_t fmcpdg = -999;
Int_t fmctype = -999;
Float_t fmceta = -999;
Float_t fmcpt = -999;

Int_t badtpc=0;

enum {
  //pp
  k10ppRAA            = 1,

  //11a has the same TPC configuration as in 2010
  k11aRAA            ,
 
  //PbPb
  k10hRAA            
};

Int_t fdata = 0;
Bool_t fkMC = 0;
Bool_t fkTEST = 0;
Int_t fkCharge = 0;
Bool_t fkPRIM = 0;

enum Cut{
  kRAA = 1,
  kOPEN = 2,
};

Long64_t Idx(const Int_t ii)
{
  return 1<<ii;
}

Float_t MomAxis()
{ 
  return TMath::Log10(fmom);
}

Float_t PtAxis(){ return TMath::Log10(fpt); }

Bool_t IsTOF(Int_t &id, const Int_t incn, const Int_t excn)
{
  const Double_t incCut = incn;
  const Double_t excCut = excn;
 
  //electron exclusion is too strong and not necessary, anyway it is very small fraction
  Bool_t kpion = kFALSE;
  Bool_t kkaon = kFALSE;
  Bool_t kproton = kFALSE;

  if(TMath::Abs(ftofpion)    <incCut && TMath::Abs(ftofkaon)>excCut && TMath::Abs(ftofproton) > excCut) kpion   = kTRUE;
  if(TMath::Abs(ftofkaon  )  <incCut && TMath::Abs(ftofpion)>excCut && TMath::Abs(ftofproton) > excCut) kkaon   = kTRUE;
  if(TMath::Abs(ftofproton)  <incCut && TMath::Abs(ftofkaon)>excCut && TMath::Abs(ftofpion)   > excCut) kproton = kTRUE;

  const Int_t finalbit = kpion+kkaon+kproton;
  if(finalbit>1){
    printf("strange IsTOF multiple!! %d %d %d\n", kpion, kkaon, kproton); exit(1);
  }

  id = kpion*dEdxGausN::kPION + kkaon*dEdxGausN::kKAON + kproton*dEdxGausN::kPROTON;

  static Int_t counter = 0;
  if(counter<10){
    printf("IsTOF incCut %e excCut %e finalbit %d id %d\n", incCut,excCut, finalbit, id); 
  }
  counter++;    


  return finalbit;
}

Bool_t DoCut()
{  
  //for secondary estimation all should be kept!
  //cut on primary for efficiency!!!
  if(fkPRIM){
    //select primary
    if(fmctype!=1)
      return kFALSE;
  }

  if(fkMC){
    //remove other particel types because their amount depends on the physics in MC (namely the generator), not always corresponding to real data
    if(fmcpdg != 11 && fmcpdg != 211 && fmcpdg != 2212 && fmcpdg != 321)
      return kFALSE;
  }

  //fkCharge=0, always true
  //fkCharge=1, fcharge<0 false, i.e. >0 true
  //fkChareg=-1, fchareg>0 false, i.e. <0 true
  if(fcharge*fkCharge<0)
    return kFALSE;

  if(!dEdxGausN::IsEtaOK(feta))
    return kFALSE;

  //PbPb
  if(fdata==Idx(k10hRAA)){
    if(!dEdxGausN::IsCenOK(fv0cen))
      return kFALSE;
  }

  Bool_t kpass = kFALSE;
  if(fdata==Idx(k10ppRAA) || fdata==Idx(k10hRAA) ||fdata==Idx(k11aRAA)){
    kpass = kTRUE;
    if(ftpcsig<EPSILON){
      if(!fkMC){
        printf("not expected!! in task pidn > 60 already made! %e \n", ftpcsig); exit(1);
      }
      else{
        kpass = kFALSE;
        badtpc++;
      }
    }
  }
  else{
    printf("DoCut wrong fdata! %d\n", fdata); exit(1);
  }

  return kpass;
}
  
void extractHist(const TString fin)
{
  //input
  TChain *ch=XGLUtils::InputFiles(fin, "tree");

  ch->SetBranchStatus("*",0);

  ch->SetBranchStatus("tpcsig",1);
  ch->SetBranchStatus("pt",1);
  ch->SetBranchStatus("mom",1);
  ch->SetBranchStatus("eta",1);
  ch->SetBranchStatus("v0cen",1);
  ch->SetBranchStatus("charge",1);
  
  ch->SetBranchAddress("tpcsig", &ftpcsig);
  ch->SetBranchAddress("pt",     &fpt);
  ch->SetBranchAddress("mom",     &fmom);
  ch->SetBranchAddress("eta",     &feta);
  ch->SetBranchAddress("v0cen",     &fv0cen);
  ch->SetBranchAddress("charge",     &fcharge);

  //ch->SetBranchAddress("tofmass",     &ftofmass);
  ch->SetBranchAddress("tofpion",     &ftofpion);
  ch->SetBranchAddress("tofkaon",     &ftofkaon);
  ch->SetBranchAddress("tofproton",     &ftofproton);
  ch->SetBranchAddress("tofelectron",     &ftofelectron);
 
  ch->SetBranchStatus("ifilled",1);
  ch->SetBranchAddress("ifilled",     &fifilled);

  if(fkMC){
    ch->SetBranchStatus("mcpdg",1);
    ch->SetBranchAddress("mcpdg",     &fmcpdg);

    ch->SetBranchStatus("mctype",1);
    ch->SetBranchAddress("mctype",     &fmctype);

    ch->SetBranchStatus("mceta",1);
    ch->SetBranchAddress("mceta",     &fmceta);

    ch->SetBranchStatus("mcpt",1);
    ch->SetBranchAddress("mcpt",     &fmcpt);
  }

  TList *ll=new TList;

  //_____________________________________________________
  //_____________________________________________________

  TH1D * hptspectrum = dEdxGausN::GetHistPtSpectrum();
  ll->Add(hptspectrum);
  TH2D * horiginal = dEdxGausN::GetHistTPCsigVsLog10P("horiginal"); 
  ll->Add(horiginal);
  //1 sigma incCut
  TH2D * h1sTOForiginal = dEdxGausN::GetHistTPCsigVsLog10P("h1sTOForiginal"); 
  ll->Add(h1sTOForiginal);
  //2 sigma incCut
  TH2D * h2sTOForiginal = dEdxGausN::GetHistTPCsigVsLog10P("h2sTOForiginal"); 
  ll->Add(h2sTOForiginal);
  //1 sigma
  TH2D *h1sTOFlogp = dEdxGausN::GetHistTOFVsLog10P("h1sTOFlogp");
  ll->Add(h1sTOFlogp);
  //2 sigma
  TH2D *h2sTOFlogp = dEdxGausN::GetHistTOFVsLog10P("h2sTOFlogp");
  ll->Add(h2sTOFlogp);

  TH2D * kk = dEdxGausN::GetHistLog10PtVsLog10P(); 
  ll->Add(kk);
 
  //only for book-keeping
  //==>
  TH1D * heta = new TH1D("heta","",200,-1,1);
  ll->Add(heta);
  TH1D * hv0cen = new TH1D("hv0cen", "", 200, 0, 100);
  ll->Add(hv0cen);
  TH1I * hcharge = new TH1I("hcharge", "", 10, -5, 5);
  ll->Add(hcharge);
 
  TH1D * hifilled = new TH1D("hifilled", "", 10, 0, 10);
  ll->Add(hifilled);

  //TH2D *htofmass = dEdxGausN::GetHistTOFMass("htofmass");
  //ll->Add(htofmass);
  //TH2D *h1stofmass = dEdxGausN::GetHistTOFMass("h1stofmass");
  //ll->Add(h1stofmass);
  //TH2D *h2stofmass = dEdxGausN::GetHistTOFMass("h2stofmass");
  //ll->Add(h2stofmass);
  
  TH2D *htofpion = dEdxGausN::GetHistTOFpar(dEdxGausN::kPION, "htofpion");
  ll->Add(htofpion);
  TH2D *htofkaon = dEdxGausN::GetHistTOFpar(dEdxGausN::kKAON, "htofkaon");
  ll->Add(htofkaon);
  TH2D *htofproton = dEdxGausN::GetHistTOFpar(dEdxGausN::kPROTON, "htofproton");
  ll->Add(htofproton);
  TH2D *htofelectron = dEdxGausN::GetHistTOFpar(dEdxGausN::kELECTRON, "htofelectron");
  ll->Add(htofelectron);
  //<==

  //MC related ============>
  TH2D * PsVrecVzhmcpdgVsPt = dEdxGausN::GetHistMCPDGVsLog10Pt("PsVrecVzhmcpdgVsPt");
  if(fkMC) ll->Add(PsVrecVzhmcpdgVsPt);

  TH2D * hmcpdgEta = new TH2D("mcpdgVsEta","",200,-1,1, 4,0,4);
  if(fkMC) ll->Add(hmcpdgEta);

  TH2D * hmcpwmPt[5];
  for(Int_t ii=0; ii<5; ii++){
    hmcpwmPt[ii]= dEdxGausN::GetHistMCPWMVsLog10Pt(ii);
    if(fkMC) ll->Add(hmcpwmPt[ii]);
  }

  TH2D * hmcptPt = dEdxGausN::GetHistMCPtVsLog10Pt();
  if(fkMC) ll->Add(hmcptPt);

  TH2D * hparsMom[4]={0x0,0x0,0x0,0x0};
  TH2D * hparsPt[4]={0x0,0x0,0x0,0x0};
  if(fkMC){
    for(Int_t ii=0; ii<4; ii++){
      hparsMom[ii] = dEdxGausN::GetHistTPCsigVsLog10P(Form("mcpar%dMom",ii));
      ll->Add(hparsMom[ii]);

      hparsPt[ii] = dEdxGausN::GetHistTPCsigVsLog10P(Form("mcpar%dPt",ii));
      ll->Add(hparsPt[ii]);
    }
  }
  //<=======================

  //original fill!
  Long64_t nent = 0;

  //only for p>20 GeV/c  ==>
  //not effective now
  Long64_t nfill = 0;
  const Long64_t nfcut = 50;
  const Double_t ptrigger = 1e9;
  //<==

  while(ch->GetEntry(nent)){
    if(nent%1000000==0){
      printf("original fill %lld\n", nent);
    }
    
    nent++;

    if(fkTEST){
      if(nent>=1000000)  break;
    }

    if(DoCut()){

      Int_t tofid = -999;
        
      //1-sigma
      if(IsTOF(tofid, 1, 4)){
        htofpion->Fill(MomAxis(), ftofpion);
        htofkaon->Fill(MomAxis(), ftofkaon);
        htofproton->Fill(MomAxis(), ftofproton);
        htofelectron->Fill(MomAxis(), ftofelectron);
        
        //h1stofmass->Fill(MomAxis(), ftofmass);
        h1sTOFlogp->Fill(MomAxis(), tofid);
      }
      else{
        h1sTOForiginal->Fill(MomAxis(), ftpcsig);        
      }
      
      //2 sigma
      if(IsTOF(tofid, 2, 4)){
        //h2stofmass->Fill(MomAxis(), ftofmass);
        h2sTOFlogp->Fill(MomAxis(), tofid);
      }
      else{
        h2sTOForiginal->Fill(MomAxis(), ftpcsig);        
      }
      
      //no TOF
      //htofmass->Fill(MomAxis(), ftofmass);
      horiginal->Fill(MomAxis(), ftpcsig);  
      
      //---      
      //be careful about X-axis definition
      //hptspectrum->Fill(PtAxis());
      hptspectrum->Fill(fpt);

      kk->Fill(MomAxis(), PtAxis());

      //---
      heta->Fill(feta);
      hv0cen->Fill(fv0cen);
      hcharge->Fill(fcharge);
      hifilled->Fill(fifilled);
 
      if(fmom > ptrigger){
        nfill++;
      }
      if(nfill > nfcut){
        break;
      }

      if(fkMC){
        const Int_t itype = dEdxGausN::PDG2Type(fmcpdg);
        if(itype>=0){
          hparsMom[itype]->Fill(MomAxis(), ftpcsig);
          hparsPt[itype]->Fill(PtAxis(), ftpcsig);
        }

        //for track in data, only kFailEvent = 0 is possible
        PsVrecVzhmcpdgVsPt->Fill(PtAxis(), itype);

        hmcpwmPt[itype]->Fill(PtAxis(), fmctype);
        hmcpwmPt[4]->Fill(PtAxis(), fmctype);
        hmcptPt->Fill(PtAxis(), TMath::Log10(fmcpt));
        hmcpdgEta->Fill(fmceta, itype);
      }
    }
  }
  printf("original fill nent: %lld\n", nent);

  //save original hyield
  TH1D *horiginalyields = dEdxGausN::GetHyield(horiginal);
  ll->Add(horiginalyields);
  TH1D *h1sTOForiginalyields = dEdxGausN::GetHyield(h1sTOForiginal);
  ll->Add(h1sTOForiginalyields);
  TH1D *h2sTOForiginalyields = dEdxGausN::GetHyield(h2sTOForiginal);
  ll->Add(h2sTOForiginalyields);

  //due to definition of Xaxis
  //ToNaturalScale(hptspectrum);

  //normalize ptSpectrum only after ToNaturalScale!!
  hptspectrum->Sumw2();
  for(Int_t ii=0; ii<=hptspectrum->GetNbinsX()+1; ii++){
    Double_t iyy = hptspectrum->GetBinContent(ii);
    if(iyy<EPSILON)
      continue;

    Double_t iee = TMath::Sqrt(iyy);

    const Double_t bw = hptspectrum->GetBinWidth(ii);
    iyy /= bw;
    iee /= bw;

    hptspectrum->SetBinContent(ii, iyy);
    hptspectrum->SetBinError(ii,iee);
  }

  //_____________________________________________________
  //_____________________________________________________

  ll->Print();
  ll->ls();

  printf("pseudo Nevent: %.0f\n", hifilled->GetBinContent(1));

  TString tag(fin);
  tag=  tag(tag.Last('/')+1,tag.Last('.')-tag.Last('/')-1);
  tag+="-";

  TString pwd = gSystem->pwd();
  pwd.ReplaceAll("/hera/alice/xlu/train/","");
  pwd.ReplaceAll("output","");
  pwd.ReplaceAll("analysis_To_hist","");
  pwd.ReplaceAll("hadron","");
  pwd.ReplaceAll("/","");

  tag+=pwd;  

  TFile *fout=new TFile(Form("histRec-%s.root", tag.Data()), "recreate");
  ll->Write();
  fout->Save();
  fout->Close();

  printf("done!!");
}

int main(int argc, char * argv[])
{
  for(int ii=0; ii<argc; ii++){
    printf("%d: %s\n", ii, argv[ii]);
  }

  if(argc!=2){
    printf("argc!=2\n"); return 1;
  }

  const TString fin(argv[1]);
  const TString pwd=gSystem->pwd();

  fdata = 0;
  if(fin.Contains("LHC10d") || fin.Contains("LHC10e") || fin.Contains("LHC10-d-e") ){
    if(pwd.Contains("RAA"))            fdata += Idx(k10ppRAA);
  }
  if(fin.Contains("LHC10h")){
    if(pwd.Contains("RAA"))            fdata += Idx(k10hRAA);
  }
  if(fin.Contains("LHC11a")){
    if(pwd.Contains("RAA"))            fdata += Idx(k11aRAA);
  }

  if(fdata==0){
    printf("error fdata!! %d %s %s\n", fdata, fin.Data(), pwd.Data());exit(1);
  }

  if(fin.Contains("MC")) fkMC = kTRUE;
  else fkMC = kFALSE;
  
  if(pwd.Contains("PRIM")) fkPRIM = kTRUE;
  else fkPRIM = kFALSE;

  if(pwd.Contains("test")) fkTEST = kTRUE;
  else fkTEST = kFALSE;

  fkCharge = dEdxGausN::GetkCharge();
 
  printf("fdata: %d; fkMC %d; fkPRIM %d; fkTEST %d; fkCharge %d;\n", fdata, fkMC, fkPRIM, fkTEST, fkCharge);

  dEdxGausN::EtaIni();
  dEdxGausN::CenIni();
  dEdxGausN::HistPrint();

  TStopwatch stwa;
  stwa.Start(kTRUE);

  extractHist(fin);

  stwa.Stop(); 
  stwa.Print("m"); //"m": milli sec; "u": micro sec
  
  printf("badtpc %d\n", badtpc);

  return 0;
}
